package com.huishu.tcm.config;

import cn.hutool.core.util.ObjectUtil;
import com.huishu.tcm.util.RateLimit;
import com.huishu.tcm.util.exception.TcmException;
import com.huishu.tcm.util.exception.TcmExceptionEnum;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.springframework.context.annotation.EnableAspectJAutoProxy;
import org.springframework.stereotype.Component;
import org.springframework.web.context.request.RequestAttributes;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;

import javax.servlet.http.HttpServletRequest;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;

@EnableAspectJAutoProxy
@Aspect
@Component
public class RateLimitAspect {
    // 使用 ConcurrentHashMap 存储限流信息
    private final Map<String, RateLimitEntry> rateLimitMap = new ConcurrentHashMap<>();

    // 内部类存储计数和截止时间
    private static class RateLimitEntry {
        int count;
        long deadline;

        RateLimitEntry(int count, long deadline) {
            this.count = count;
            this.deadline = deadline;
        }
    }

    @Around("@annotation(rateLimit)")
    public  Object around(ProceedingJoinPoint pjp, RateLimit rateLimit) throws Throwable {
        HttpServletRequest request = ((ServletRequestAttributes) RequestContextHolder.getRequestAttributes()).getRequest();
        String ip = request.getRemoteAddr();
        String key = "rateLimit" + ip; // 只用 IP 作为 key
        int limit = rateLimit.limit(); // 应为 1
        int time = rateLimit.time();   // 应为 20 秒

        long now = new Date().getTime();
        RateLimitEntry rateLimitEntry = rateLimitMap.get(key);

        // 检查现有记录
        if (ObjectUtil.isNotEmpty(rateLimitEntry)) {
            long deadline = rateLimitEntry.deadline;
            if (now < deadline) {
                // 未过期，检查是否达到限制
                if (rateLimitEntry.count >= limit) {
                    throw new TcmException(TcmExceptionEnum.SERVICE_TOO_MANY_REQUESTS);
                }
                // 未达到限制，计数加 1
                rateLimitEntry.count += 1;
                rateLimitMap.put(key, rateLimitEntry);
            } else {
                // 已过期，移除旧记录并创建新记录
                rateLimitMap.remove(key);
                rateLimitEntry = new RateLimitEntry(1, now + time * 1000);
                rateLimitMap.put(key, rateLimitEntry);
            }
        } else {
            // 无记录，创建新记录
            rateLimitEntry = new RateLimitEntry(1, now + time * 1000);
            rateLimitMap.put(key, rateLimitEntry);
        }
        return pjp.proceed();
    }
}
